package my.shadowsocks.common.encryption.impl;

import my.shadowsocks.common.encryption.AbstractCrypt;
import org.bouncycastle.crypto.StreamBlockCipher;
import org.bouncycastle.crypto.engines.AESEngine;
import org.bouncycastle.crypto.modes.CFBBlockCipher;
import org.bouncycastle.crypto.modes.OFBBlockCipher;

import java.io.ByteArrayOutputStream;
import java.security.InvalidAlgorithmParameterException;
import java.util.HashMap;
import java.util.Map;


/**
 * AES 实现类
 * @author fengbo
 */
public class AesCrypt extends AbstractCrypt {

	private final static String CIPHER_AES_128_CFB = "aes-128-cfb";
	private final static String CIPHER_AES_192_CFB = "aes-192-cfb";
	private final static String CIPHER_AES_256_CFB = "aes-256-cfb";
	private final static String CIPHER_AES_128_OFB = "aes-128-ofb";
	private final static String CIPHER_AES_192_OFB = "aes-192-ofb";
	private final static String CIPHER_AES_256_OFB = "aes-256-ofb";

	public static Map<String, String> getCiphers() {
		Map<String, String> ciphers = new HashMap<>(10);
		ciphers.put(CIPHER_AES_128_CFB, AesCrypt.class.getName());
		ciphers.put(CIPHER_AES_192_CFB, AesCrypt.class.getName());
		ciphers.put(CIPHER_AES_256_CFB, AesCrypt.class.getName());
		ciphers.put(CIPHER_AES_128_OFB, AesCrypt.class.getName());
		ciphers.put(CIPHER_AES_192_OFB, AesCrypt.class.getName());
		ciphers.put(CIPHER_AES_256_OFB, AesCrypt.class.getName());
		return ciphers;
	}

	public AesCrypt(String name, String password) {
		super(name, password);
	}

	@Override
	public int getKeyLength() {
		switch (name) {
			case CIPHER_AES_128_CFB:
				return 16;
			case CIPHER_AES_192_CFB:
			case CIPHER_AES_192_OFB:
				return 24;
			case CIPHER_AES_256_CFB:
			case CIPHER_AES_256_OFB:
				return 32;
			default:
				return 0;
		}
	}

	@Override
	protected StreamBlockCipher getCipher(boolean isEncrypted)
			throws InvalidAlgorithmParameterException {
		AESEngine engine = new AESEngine();
		switch (name) {
			case CIPHER_AES_128_CFB:
				return new CFBBlockCipher(engine, getIVLength() * 8);
			case CIPHER_AES_192_CFB:
				return new CFBBlockCipher(engine, getIVLength() * 8);
			case CIPHER_AES_256_CFB:
				return new CFBBlockCipher(engine, getIVLength() * 8);
			case CIPHER_AES_128_OFB:
				return new OFBBlockCipher(engine, getIVLength() * 8);
			case CIPHER_AES_192_OFB:
				return new OFBBlockCipher(engine, getIVLength() * 8);
			case CIPHER_AES_256_OFB:
				return new OFBBlockCipher(engine, getIVLength() * 8);
			default:
				throw new InvalidAlgorithmParameterException(name);
		}
	}

	@Override
	public int getIVLength() {
		return 16;
	}

	@Override
	protected void encrypt0(byte[] data, ByteArrayOutputStream stream) {
		int noBytesProcessed;
		byte[] buffer = new byte[data.length];

		noBytesProcessed = encCipher.processBytes(data, 0, data.length, buffer, 0);
		stream.write(buffer, 0, noBytesProcessed);
	}

	@Override
	protected void decrypt0(byte[] data, ByteArrayOutputStream stream) {
		int noBytesProcessed;
		byte[] buffer = new byte[data.length];

		noBytesProcessed = decCipher.processBytes(data, 0, data.length, buffer, 0);
		stream.write(buffer, 0, noBytesProcessed);
	}
}
